#!/usr/bin/env python3
"""
阶段2：人工标注角色
交互式标注说话人身份
输入：raw_transcript.md
输出：annotated_transcript.md
"""

import json
import sys
from pathlib import Path
from typing import Dict, List
import logging

logger = logging.getLogger(__name__)

class Stage2Processor:
    """人工标注处理器"""
    
    def __init__(self, stage1_data: Dict, output_dir: Path):
        self.stage1_data = stage1_data
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
    def get_speaker_roles(self) -> Dict[str, str]:
        """获取说话人角色映射"""
        speakers = set(segment["speaker"] for segment in self.stage1_data["data"]["segments"])
        
        print("\n" + "="*50)
        print("阶段2：人工标注说话人角色")
        print("="*50)
        
        role_mapping = {}
        role_options = {
            "1": "需求方（客户）",
            "2": "需求方（业务人员）", 
            "3": "乙方（技术）",
            "4": "乙方（产品）",
            "5": "乙方（项目经理）",
            "6": "其他"
        }
        
        for speaker in sorted(speakers):
            print(f"\n请为 {speaker} 选择角色：")
            for key, value in role_options.items():
                print(f"  {key}. {value}")
            
            while True:
                choice = input(f"选择 {speaker} 的角色 (1-6): ").strip()
                if choice in role_options:
                    role_mapping[speaker] = role_options[choice]
                    break
                else:
                    print("请输入 1-6 之间的数字")
        
        return role_mapping
    
    def process(self) -> Dict:
        """执行人工标注"""
        logger.info("阶段2：开始人工标注")
        
        # 获取角色映射
        role_mapping = self.get_speaker_roles()
        
        # 更新数据
        annotated_data = self.stage1_data.copy()
        annotated_data["role_mapping"] = role_mapping
        
        # 更新每个段落的角色信息
        for segment in annotated_data["data"]["segments"]:
            original_speaker = segment["speaker"]
            segment["role"] = role_mapping.get(original_speaker, "未知")
        
        # 保存标注转录文件
        output_file = self.output_dir / "stage2" / "annotated_transcript.md"
        output_file.parent.mkdir(parents=True, exist_ok=True)
        
        self._save_annotated_transcript(annotated_data, output_file)
        
        logger.info(f"阶段2完成：{output_file}")
        return {
            "stage2_file": str(output_file),
            "data": annotated_data,
            "role_mapping": role_mapping
        }
    
    def _save_annotated_transcript(self, data: Dict, output_file: Path):
        """保存标注转录文件"""
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write("# 标注转录文件\n\n")
            f.write("**角色映射**:\n")
            for speaker, role in data["role_mapping"].items():
                f.write(f"- {speaker}: {role}\n")
            f.write("\n---\n\n")
            
            for segment in data["data"]["segments"]:
                f.write(f"## {segment['role']} [{segment['timestamp']}]\n")
                f.write(f"原始说话人: {segment['speaker']}\n")
                f.write(f"{segment['text']}\n\n")
        
        # 保存JSON格式
        json_file = output_file.with_suffix('.json')
        with open(json_file, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)

if __name__ == "__main__":
    import sys
    if len(sys.argv) != 2:
        print("用法: python stage2_annotate.py <stage1_json_file>")
        sys.exit(1)
    
    with open(sys.argv[1], 'r', encoding='utf-8') as f:
        stage1_data = json.load(f)
    
    processor = Stage2Processor(stage1_data, "../output")
    result = processor.process()
    print(f"处理完成: {result['stage2_file']}")